"""Provides an automatically-cleaned temporary directory.

All functions in this module are thread-safe.
"""

import glob
import os
import shutil
import sys
import tempfile

TEMP_PID_FILE = '.expect_tests_pidfile'
TEMP_SUFFIX = '.expect_tests_temp'


class TempDir(object):
  """Context manager for a temporary directory.

  Use tempfile.TemporaryDirectory instead with Python 3.
  """
  def __init__(self, leak=False):
    self.old_temp_dir = tempfile.tempdir
    self.leak = leak

  def __enter__(self):
    tempfile.tempdir = set_temp_dir()
    if self.leak:
      print "leaking temporary directory:", tempfile.tempdir

  def __exit__(self, _type, _value, _traceback):
    try:
      if not self.leak:
        clear_temp_dir(tempfile.tempdir)
    finally:
      tempfile.tempdir = self.old_temp_dir


def set_temp_dir(suffix=TEMP_SUFFIX):
  """Provides a temporary directory for tests.

  This function also takes care of removing old temp directories from dead
  processes.
  """
  temp_dir = tempfile.mkdtemp(suffix)

  # Clean up obsolete directories.
  for p in glob.glob(os.path.join(os.path.dirname(temp_dir), '*'+suffix)):
    if p == temp_dir:
      continue

    pfile = os.path.join(p, TEMP_PID_FILE)
    if os.path.exists(pfile):
      with open(pfile, 'rb') as f:
        try:
          os.kill(int(f.read()), 0)
          continue
        except (OSError, TypeError):
          pass

    shutil.rmtree(p, ignore_errors=True)

  # Create current PID file.
  with open(os.path.join(temp_dir, TEMP_PID_FILE), 'wb') as f:
    f.write(str(os.getpid()))
  return temp_dir


def clear_temp_dir(temp_dir):
  # Try to nuke the pidfile first.
  pfile = os.path.join(temp_dir, TEMP_PID_FILE)
  try:
    os.unlink(pfile)
  except OSError as e:
    print >> sys.stderr, "Error removing %r: %s" % (pfile, e)
  shutil.rmtree(temp_dir, ignore_errors=True)
